import matplotlib.pyplot as plt
import numpy as np
import util

from p05b_lwr import LocallyWeightedLinearRegression


def main(tau_values, train_path, valid_path, test_path, pred_path):
    """Problem 5(b): Tune the bandwidth paramater tau for LWR.

    Args:
        tau_values: List of tau values to try.
        train_path: Path to CSV file containing training set.
        valid_path: Path to CSV file containing validation set.
        test_path: Path to CSV file containing test set.
        pred_path: Path to save predictions.
    """
    # Load training set
    x_train, y_train = util.load_dataset(train_path, add_intercept=True)

    # *** START CODE HERE ***
    x_valid, y_valid = util.load_dataset(valid_path, add_intercept=True)

    # Search tau_values for the best tau (lowest MSE on the validation set)
    best_tau = None
    best_mse = float('inf')
    best_model = None
    mse_values = []
    all_predictions = []

    print("Tuning tau parameter:")
    for tau in tau_values:
    # Fit a LWR model with the best tau value
        model_c = LocallyWeightedLinearRegression(tau)
        model_c.fit(x_train, y_train)

        y_pred = model_c.predict(x_valid)
        all_predictions.append(y_pred)

        mse = np.mean((y_pred - y_valid)** 2)
        mse_values.append(mse)
        print(f"τ = {tau:.2f}, MSE = {mse:.4f}")

        if mse < best_mse:
            best_tau = tau
            best_mse = mse
            best_model = model_c

    print(f"\nBest τ = {best_tau:.2f} with MSE = {best_mse:.4f}")
    # Load test set
    x_test, y_test = util.load_dataset(test_path, add_intercept=True)

    # Run on the test set to get the MSE value
    y_test_pred = best_model.predict(x_test)
    test_mse = np.mean((y_test_pred - y_test) ** 2)

    print(f"Test MSE with best τ = {test_mse:.4f}")

    np.savetxt(pred_path, y_test_pred)

    # Plot data for each tau value (使用已保存的预测结果)
    plt.figure(figsize=(15, 10))

    for i, tau in enumerate(tau_values):
        plt.subplot(2, 3, i + 1)

        y_pred = all_predictions[i]
        # Plot training data
        plt.plot(x_train[:, 1], y_train, 'bx', markersize=4, label='Training Data')

        # Plot validation predictions
        plt.plot(x_valid[:, 1], y_pred, 'ro', markersize=4, label='Validation Predictions')

        plt.xlabel('x')
        plt.ylabel('y')
        plt.title(f'τ = {tau:.2f}, MSE = {mse_values[i]:.4f}')
        plt.legend()
        plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig('output/p05c_tau_tuning.png', dpi=300, bbox_inches='tight')
    plt.show()
    # *** END CODE HERE ***

    # Run on the test set to get the MSE value
    # Save predictions to pred_path
    # Plot data
    # *** END CODE HERE ***
